import uuid
from typing import TypedDict, Optional

from langgraph.graph import StateGraph
from langgraph.constants import START, END
from langgraph.types import interrupt, Command
from langgraph.checkpoint.memory import InMemorySaver


class State(TypedDict):
    """The graph state."""
    age: Optional[str]
    name: Optional[str]


def human_node(state: State):
    # 多个 interrupt，严格按顺序
    name = state["name"]
    age = state["age"]

    if name is None:
        name = interrupt("what is your name?")
    if age is None:
        age = interrupt("what is your age?")

    print(f"Name: {name}, Age: {age}")
    return {
        "name": name,
        "age": age,
    }


if __name__ == '__main__':
    builder = StateGraph(State)
    builder.add_node("human_node", human_node)
    builder.set_entry_point("human_node")
    builder.set_finish_point("human_node")

    checkpointer = InMemorySaver()
    graph = builder.compile(checkpointer=checkpointer)

    config = {
        "configurable": {
            "thread_id": str(uuid.uuid4()),
        }
    }

    # 第一次运行：触发 name 中断
    for chunk in graph.stream({"age": None, "name": None}, config):
        print("[第一次运行]:", chunk)

    # 回答 name
    for chunk in graph.stream(Command(resume="Alice"), config):
        print("[第二次运行 - 回答 name]:", chunk)

    # 回答 age
    for chunk in graph.stream(Command(resume="25"), config):
        print("[第三次运行 - 回答 age]:", chunk)

